// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/prestosql/aggregates/array_agg_aggregate.h>
#include <pollux/exec/container_row_serde.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/aggregates/value_list.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        struct ArrayAccumulator {
            ValueList elements;
        };

        class ArrayAggAggregate : public exec::Aggregate {
        public:
            explicit ArrayAggAggregate(TypePtr resultType, bool ignoreNulls)
                : Aggregate(resultType), ignoreNulls_(ignoreNulls) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(ArrayAccumulator);
            }

            bool isFixedSize() const override {
                return false;
            }

            bool supportsToIntermediate() const override {
                return true;
            }

            void toIntermediate(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                VectorPtr &result) const override {
                const auto &elements = args[0];

                const auto numRows = rows.size();

                // Convert input to a single-entry array.

                // Set nulls for rows not present in 'rows'.
                auto *pool = allocator_->pool();
                BufferPtr nulls = allocate_nulls(numRows, pool);
                auto mutable_nulls = nulls->asMutable<uint64_t>();
                memcpy(
                    nulls->asMutable<uint64_t>(),
                    rows.asRange().bits(),
                    bits::nbytes(numRows));

                auto loadedElements = BaseVector::loaded_vector_shared(elements);

                if (ignoreNulls_ && loadedElements->may_have_nulls()) {
                    rows.applyToSelected([&](vector_size_t row) {
                        if (loadedElements->is_null_at(row)) {
                            bits::set_null(mutable_nulls, row);
                        }
                    });
                }

                // Set offsets to 0, 1, 2, 3...
                BufferPtr offsets = allocateOffsets(numRows, pool);
                auto *rawOffsets = offsets->asMutable<vector_size_t>();
                std::iota(rawOffsets, rawOffsets + numRows, 0);

                // Set sizes to 1.
                BufferPtr sizes = allocateSizes(numRows, pool);
                auto *rawSizes = sizes->asMutable<vector_size_t>();
                std::fill(rawSizes, rawSizes + numRows, 1);

                result = std::make_shared<ArrayVector>(
                    pool,
                    ARRAY(elements->type()),
                    nulls,
                    numRows,
                    offsets,
                    sizes,
                    loadedElements);
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto vector = (*result)->as<ArrayVector>();
                POLLUX_CHECK(vector);
                vector->resize(numGroups);

                auto elements = vector->elements();
                elements->resize(countElements(groups, numGroups));

                uint64_t *rawNulls = getRawNulls(vector);
                vector_size_t offset = 0;
                for (int32_t i = 0; i < numGroups; ++i) {
                    auto &values = value<ArrayAccumulator>(groups[i])->elements;
                    auto arraySize = values.size();
                    if (arraySize) {
                        clearNull(rawNulls, i);

                        ValueListReader reader(values);
                        for (auto index = 0; index < arraySize; ++index) {
                            reader.next(*elements, offset + index);
                        }
                        vector->setOffsetAndSize(i, offset, arraySize);
                        offset += arraySize;
                    } else {
                        vector->set_null(i, true);
                    }
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                extractValues(groups, numGroups, result);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodedElements_.decode(*args[0], rows);
                rows.applyToSelected([&](vector_size_t row) {
                    if (ignoreNulls_ && decodedElements_.is_null_at(row)) {
                        return;
                    }
                    auto group = groups[row];
                    auto tracker = trackRowSize(group);
                    value<ArrayAccumulator>(group)->elements.appendValue(
                        decodedElements_, row, allocator_);
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                decodedIntermediate_.decode(*args[0], rows);

                auto array_vector = decodedIntermediate_.base()->as<ArrayVector>();
                auto &elements = array_vector->elements();
                rows.applyToSelected([&](vector_size_t row) {
                    auto group = groups[row];
                    auto decodedRow = decodedIntermediate_.index(row);
                    auto tracker = trackRowSize(group);
                    if (!decodedIntermediate_.is_null_at(row)) {
                        value<ArrayAccumulator>(group)->elements.appendRange(
                            elements,
                            array_vector->offsetAt(decodedRow),
                            array_vector->sizeAt(decodedRow),
                            allocator_);
                    }
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /* mayPushdown */) override {
                auto &values = value<ArrayAccumulator>(group)->elements;

                decodedElements_.decode(*args[0], rows);
                auto tracker = trackRowSize(group);
                rows.applyToSelected([&](vector_size_t row) {
                    if (ignoreNulls_ && decodedElements_.is_null_at(row)) {
                        return;
                    }
                    values.appendValue(decodedElements_, row, allocator_);
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /* mayPushdown */) override {
                decodedIntermediate_.decode(*args[0], rows);
                auto array_vector = decodedIntermediate_.base()->as<ArrayVector>();

                auto &values = value<ArrayAccumulator>(group)->elements;
                auto elements = array_vector->elements();
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decodedIntermediate_.is_null_at(row)) {
                        auto decodedRow = decodedIntermediate_.index(row);
                        values.appendRange(
                            elements,
                            array_vector->offsetAt(decodedRow),
                            array_vector->sizeAt(decodedRow),
                            allocator_);
                    }
                });
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                for (auto index: indices) {
                    new(groups[index] + offset_) ArrayAccumulator();
                }
            }

            void destroyInternal(melon::Range<char **> groups) override {
                for (auto group: groups) {
                    if (isInitialized(group)) {
                        value<ArrayAccumulator>(group)->elements.free(allocator_);
                    }
                }
            }

        private:
            vector_size_t countElements(char **groups, int32_t numGroups) const {
                vector_size_t size = 0;
                for (int32_t i = 0; i < numGroups; ++i) {
                    size += value<ArrayAccumulator>(groups[i])->elements.size();
                }
                return size;
            }

            // A boolean representing whether to ignore nulls when aggregating inputs.
            const bool ignoreNulls_;
            // Reusable instance of DecodedVector for decoding input vectors.
            DecodedVector decodedElements_;
            DecodedVector decodedIntermediate_;
        };
    } // namespace

    void registerArrayAggAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("E")
            .returnType("array(E)")
            .intermediateType("array(E)")
            .argumentType("E")
            .build()
        };

        auto name = prefix + kArrayAgg;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig &config) -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(
                    argTypes.size(), 1, "{} takes at most one argument", name);
                return std::make_unique<ArrayAggAggregate>(
                    resultType, config.prestoArrayAggIgnoreNulls());
            },
            withCompanionFunctions,
            overwrite);
    }

    void registerInternalArrayAggAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("E")
            .returnType("array(E)")
            .intermediateType("array(E)")
            .argumentType("E")
            .build()
        };

        auto name = prefix + "$internal$array_agg";
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step /*step*/,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(
                    argTypes.size(), 1, "{} takes at most one argument", name);
                return std::make_unique<ArrayAggAggregate>(
                    resultType, /*ignoreNulls*/ false);
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
